import numpy as np
import numpy.random as rand
import math
from tqdm import tqdm
import config_ql as qc
import qlagent
from helpers.data_saver import DataSaver

N = qc.beacon_num
#设置8个信标的位置
def set_beacon():
    beacon_pos = np.array([[0, 0]])
    for i in range(N):
        beacon_pos = np.insert(beacon_pos, i, np.array([rand.randint(0, qc.boundary), rand.randint(0, qc.boundary)]),
                               axis=0)
    beacon_pos = np.delete(beacon_pos, len(beacon_pos) - 1, axis=0)
    #print('beacon_pos', beacon_pos)
    return beacon_pos

#设置目标的位置
def set_target():
    target_pos = np.array([rand.randint(0, qc.boundary), rand.randint(0, qc.boundary)])
    #h = rand.randint(0, qc.boundary)
    h = 600
    #print('target_pos:', target_pos)
    return target_pos, h

#距离暂时用实际距离代替，因为不清楚论文中的Ti(k)如何求
def distance(beacon_pos, target_pos):
    dist = np.zeros(N)
    for i in range(N):
        dist[i] = np.sqrt((beacon_pos[i][0] - target_pos[0])**2 + (beacon_pos[i][1] - target_pos[1])**2)
    #print('dist:', dist)
    return dist

#构造论文的A矢量
def construct_A(action, select_num, beacon_pos):
    beacon_max = 0  # 记录最大的信标
    beacons_num = 0  # 记录选择的信标数
    for i in range(8):
        if action >= 2**(7-i):
            beacon_max = 7-i
            break
    for i in select_num:
        if i == 1:
            beacons_num += 1
    if beacons_num < 3:
        beacon_max = 2
        select_num = np.array([1, 1, 1, 0, 0, 0, 0, 0])
    A = np.array([[0, 0]])
    num_1 = 0  # 用来计数所有信标数
    num_2 = 0  # 用来计数被选择的坐标
    for i in select_num:
        if i == 1:
            A = np.insert(A, num_2, beacon_pos[num_1] - beacon_pos[beacon_max], axis=0)
            num_2 += 1
        num_1 += 1
    A = np.delete(A, len(A) - 1, axis=0)
    A = np.delete(A, len(A) - 1, axis=0)
    #print('A:', A)
    return A

#构造论文的b矢量
def construct_b(action, select_num, beacon_pos, dist):
    beacon_max = 0  # 记录最大的信标
    beacons_num = 0  # 记录选择的信标数
    for i in range(8):
        if action >= 2**(7-i):
            beacon_max = 7-i
            break
    for i in select_num:
        if i == 1:
            beacons_num += 1
    if beacons_num < 3:
        beacon_max = 2
        select_num = np.array([1, 1, 1, 0, 0, 0, 0, 0])
    b = np.array([[0]])
    num_1 = 0  # 用来计数所有信标数
    num_2 = 0  # 用来计数被选择的坐标
    for i in select_num:
        if i == 1:
            a = beacon_pos[num_1][0]**2 + beacon_pos[num_1][1]**2 - beacon_pos[beacon_max][0]**2\
                - beacon_pos[beacon_max][1]**2 + dist[beacon_max]**2 - dist[num_1]**2
            b = np.insert(b, num_2, a, axis=0)
            num_2 += 1
        num_1 += 1
    b = np.delete(b, len(b) - 1, axis=0)
    b = np.delete(b, len(b) - 1, axis=0)
    #print('b:', b)
    return b

# 计算目标的位置
def target_position(A, b):
    x = np.dot(A.T, A)
    if x[0][1] == x[1][1] == x[1][0] == x[0][0] == 0:
        x = np.array([[0.0000001, 0.0000002], [0.0000003, 0.0000004]])
    if x[0][1] == x[1][1] == x[1][0] == x[0][0]:
        x = np.array([[x[0][0]+0.0000001, x[0][1]+0.0000002], [x[1][0]+0.0000003, x[1][1]+0.0000004]])
    x = np.matrix(x)
    y = np.dot(x.I, A.T)
    l = 0.5 * np.dot(y, b)
    #print('l:', l)
    return l

# 计算RMSE
def Localization_error(l, dist, h, beacon_pos, select_num):
    num_1 = 0  # 用来计数所有信标数
    num_2 = 0  # 用来计数被选择的坐标
    sum = 0
    beacon_max = 0  # 记录最大的信标
    beacons_num = 0  # 记录选择的信标数
    for i in select_num:
        if i == 1:
            beacons_num += 1
    if beacons_num < 3:
        beacon_max = 2
        select_num = np.array([1, 1, 1, 0, 0, 0, 0, 0])
    for i in select_num:
        if i == 1:
            sum += abs(dist[num_1]**2 - (l[0] - beacon_pos[num_1][0])**2 - (l[1] - beacon_pos[num_1][1])**2)
            num_2 += 1
        num_1 += 1
    r = (1/(num_2+1)) * sum
    #print('r:', r)
    return r

# 计算Energy Consumption
def Energy_Consumption(select_num):
    Energy_Consumption = 0
    beacon_max = 0  # 记录最大的信标
    beacons_num = 0  # 记录选择的信标数
    for i in select_num:
        if i == 1:
            beacons_num += 1
    if beacons_num < 3:
        beacon_max = 2
        select_num = np.array([1, 1, 1, 0, 0, 0, 0, 0])
    for i in select_num:
        if i == 1:
            Energy_Consumption += qc.beacon_transmit_power
    #print('Energy_Consumption', Energy_Consumption, 'select_num', select_num)
    return Energy_Consumption

def main():
    saver = DataSaver('RL')
    Epsilon = {'epsi_low': 0.01, 'epsi_high': 0.9, 'decay': 500}
    agent = qlagent.QLAgent()
    beacon_pos = set_beacon()
    target_pos, h = set_target()
    for episode in range(qc.EPISODE_NUMBER):
        print(episode)
        state_results = qc.reset()
        r = state_results['Localization_Error']
        p = state_results['beacons'] * qc.beacon_transmit_power
        utility = qc.reward_function(r, p)
        saver.save_step(episode, 0, utility, state_results)
        for time_slot in tqdm(range(qc.STEP_NUMBER)):
            epsilon = Epsilon['epsi_low'] + (Epsilon['epsi_high'] - Epsilon['epsi_low']) * (
                        math.exp(-2.0 * time_slot / Epsilon['decay']))
            action_id, select_num = agent.choose_action(state_results, epsilon)
            dist = distance(beacon_pos, target_pos)
            A = construct_A(action_id, select_num, beacon_pos)
            b = construct_b(action_id, select_num, beacon_pos, dist)
            l = target_position(A, b)
            #print(l, target_pos)
            r = Localization_error(l, dist, h, beacon_pos, select_num)
            p = Energy_Consumption(select_num)
            utility = qc.reward_function(r, p)
            state_results_next = {
                'Depth': state_results['Depth'],
                'beacons': p / 2,
                'Localization_Error': r,
            }
            agent.learn(state_results, state_results_next, action_id, utility + 100)
            saver.save_step(episode, time_slot, utility, state_results_next)
            state_results = state_results_next
    saver.save_mat()
    saver.plot_mat()

if __name__ == '__main__':
    main()